import os

from langchain.chains.combine_documents.stuff import StuffDocumentsChain
from langchain.chains.llm import LLMChain
from langchain.chains.summarize import load_summarize_chain
from langchain_community.document_loaders import WebBaseLoader
from langchain_core.prompts import PromptTemplate
from langchain_openai import ChatOpenAI
from pydantic import SecretStr

# 读取API密钥
api_key = os.getenv("DASHSCOPE_API_KEY")
if not api_key:
    raise ValueError("请设置环境变量DASHSCOPE_API_KEY（阿里云百炼API-KEY）")

# 创建大语言模型实例
model = ChatOpenAI(
    model="qwen-plus-latest",
    temperature=0.5,
    max_tokens=None,
    timeout=None,
    max_retries=2,
    api_key=SecretStr(api_key),
    base_url="https://dashscope.aliyuncs.com/compatible-mode/v1",
)

# 加载文档,使用langchain提供的WebBaseModel进行加载网上博文
loader = WebBaseLoader("https://lilianweng.github.io/posts/2023-06-23-agent/")
docs = loader.load()  # 得到篇文章

# # 第一种 Stuff
# chain = load_summarize_chain(model, chain_type='stuff')
# result = chain.invoke(docs)
# print(result['output_text'])

# 第二种：回复中文
prompt_template = """
    针对下面的内容，写一个简洁的总结摘要:
    "{text}"
    简洁的总结摘要:
"""
prompt = PromptTemplate.from_template(prompt_template)

llm_chain = LLMChain(llm=model, prompt=prompt)

stuff_chain = StuffDocumentsChain(llm_chain=llm_chain, document_variable_name='text')

result = stuff_chain.invoke(docs)
print(result['output_text'])
